# relate modules to traits

# Rscript step3.R --exprFile "../test_data/datExpr.csv" --traitsfile "../test_data/datTraits.csv" --network "../1_construct_network/network.RData" --output "../3_relate_modules_to_traits"

library(optparse)


option_list = list(
  make_option("--exprFile", type="character", default=NULL,
              help="expression file path"),
  make_option("--traitFile", type="character", default=NULL,
              help="traits file path"),
  make_option("--network", type="character", default=NULL,
              help="network file path"),
  make_option("--output", type="character", default="out.txt",
              help="output path [default= %default]")
)

# 解析命令行参数
opt_parser = OptionParser(option_list=option_list, add_help_option=TRUE)
opts = parse_args(opt_parser)

datExpr = read.csv(opts$exprFile, stringsAsFactors = FALSE, row.names = 1);
datTraits = read.csv(opts$traitFile, stringsAsFactors = FALSE, row.names = 1);


library(WGCNA)

load(opts$network)
nSamples = length(net$goodSamples)

# Recalculate MEs with color labels
moduleColors = labels2colors(net$colors)
MEs0 = moduleEigengenes(datExpr, moduleColors)$eigengenes
MEs = orderMEs(MEs0)
moduleTraitCor = cor(MEs, datTraits, use = "p");
moduleTraitPvalue = corPvalueStudent(moduleTraitCor, nSamples);
textMatrix = paste(signif(moduleTraitCor, 2), "\n(",
                   signif(moduleTraitPvalue, 1), ")", sep = "");

png(
  filename = paste(opts$output, "module_trait_heatmap.png", sep="/"),
  width = 13,
  height = 9,
  units = "in",
  res = 300
)

par(mar = c(6, 8.5, 3, 3));
labeledHeatmap(Matrix = moduleTraitCor,
               xLabels = names(datTraits),
               yLabels = names(MEs),
               ySymbols = names(MEs),
               colorLabels = FALSE,
               colors = blueWhiteRed(50),
               textMatrix = textMatrix,
               setStdMargins = FALSE,
               cex.text = 0.5,
               zlim = c(-1,1),
               main = paste("Module-trait relationships"))
dev.off()


# 计算FDR校正后的P值
moduleTraitFDR = matrix(p.adjust(moduleTraitPvalue, method = "BH"),
                       nrow = nrow(moduleTraitPvalue),
                       dimnames = dimnames(moduleTraitPvalue))

# 创建完整的结果数据框
result_df <- data.frame(
    Module = rep(rownames(moduleTraitPvalue), ncol(moduleTraitPvalue)),
    ModuleColor = rep(moduleColors[match(rownames(moduleTraitPvalue), names(MEs))],
                     ncol(moduleTraitPvalue)),
    Trait = rep(colnames(moduleTraitPvalue), each = nrow(moduleTraitPvalue)),
    Correlation = as.vector(moduleTraitCor),
    Pvalue = as.vector(moduleTraitPvalue),
    FDR = as.vector(moduleTraitFDR)
)

# 按P值排序
result_df <- result_df[order(result_df$Pvalue),]

# 保存完整结果
write.csv(result_df,
          file = paste(opts$output, "module_trait_heatmap.csv", sep="/"),
          row.names = FALSE)

# sample trait cluster
sampleTree = hclust(dist(datExpr), method = "average");
traitColors = numbers2colors(datTraits, signed = FALSE);

png(
  filename = paste(opts$output, "sample_trait_cluster.csv", sep="/"),
  width = 13,
  height = 9,
  units = "in",
  res = 300
)

plotDendroAndColors(sampleTree, traitColors,
                    groupLabels = names(datTraits),
                    main = "Sample dendrogram and trait heatmap")

dev.off()
